#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/10/30 8:56 上午
# @Author  : WangZhixing

from typing import Optional, Callable, List
import torch
from torch_geometric.data import Data,InMemoryDataset
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import OneHotEncoder
import random
from .DataInFo.DataInfo import DataInfo
from .Process.Processrmgr import Process


def index_to_mask(index, size):
    mask = torch.zeros((size,), dtype=torch.bool)
    mask[index] = 1
    return mask


class DependenceGraph_rmgr(InMemoryDataset):
    def __init__(self, root: str,
                 transform: Optional[Callable] = None,
                 pre_transform: Optional[Callable] = None):
        super().__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[3])
        DataInfo(self.data)

    # 存放在root的目录下的raw文件夹中，如果没有调用下载
    @property
    def raw_file_names(self) -> List[str]:
        return ['edge.rsf', 'groundtruth.rsf']

    # 储存处理后的文件列表，若全部找到就跳过预处理。
    @property
    def processed_file_names(self) -> List[str]:
        return ['idx.edge.csv', 'idx.name.csv', 'idx.label.csv', 'data.pt','ground_truth.rsf']

    # 处理原始数据，并保存到processed_dir
    def process(self):
        with open(self.raw_paths[0], 'r') as f:
            process = Process()
            # 读取节点的Label，这个作用是提取label转化为数字1，2，3，4，5，6。
            # 0代表的是ground-truth中没有的节点。
            # 将节点的边，节点的名字，节点的label保存下来。
            process.Node2Num(self.raw_paths[0],
                                 self.processed_paths[0],
                                 self.processed_paths[1]
                                 )

            process.NewGroundTruth(self.raw_paths[1], self.processed_paths[4])
            process.Label2File(self.processed_paths[2])
            # 节点的pairs，保存到edge_index
            edge_index = [[lst[0], lst[1]] for lst in process.lst]
            edge_index = torch.tensor(edge_index, dtype=torch.long)
            edge_index = edge_index - edge_index.min()
            edge_index = edge_index.t().contiguous()
            # 保存边的权重
            edge_attr = [lst[2] for lst in process.lst]
            edge_attr = torch.tensor(edge_attr, dtype=torch.float)
            # 保存节点的label
            label = [lst for lst in process.NodeLabel.values()]
            label = torch.tensor(label, dtype=torch.long)

            # 节点使用one-hot编码
            data = [lst for lst in process.dic.values()]
            label_encoder = LabelEncoder()
            integer_encoded = label_encoder.fit_transform(data)
            onehot_encoder = OneHotEncoder(sparse=False)
            integer_encoded = integer_encoded.reshape(len(integer_encoded), 1)
            onehot_encoded = onehot_encoder.fit_transform(integer_encoded)
            x = torch.from_numpy(onehot_encoded).to(torch.float)

            # 划分节点
            mask = [i for i in range(process.idx)]
            random.shuffle(mask)

            idx_train = mask[0:int(len(mask)*20 / 100)]

            # idx_val = mask[int(len(mask) *1/ 10):int(3 * len(mask) *9/ 10)]

            idx_test = mask[int(20 * len(mask) / 100):]

            train_mask = index_to_mask(idx_train, len(mask))

            # val_mask = index_to_mask(idx_val, len(mask))

            test_mask = index_to_mask(idx_test, len(mask))
            #保存到data中
            data = Data(x=x, edge_index=edge_index, y=label, edge_attr=edge_attr, train_mask=train_mask
                        ,test_mask=test_mask
                        # , val_mask=val_mask
                        )
            # 是否归一化
            if self.pre_transform is not None:
                data = self.pre_transform(data)
            # 数据，保存到pt文件中，方便下次读取
            torch.save(self.collate([data]), self.processed_paths[3])
            self.dic = process.dic
            # 打印data数据信息


#
# if __name__ == '__main__':
#     c = BashFileDep(root="/Users/wzx/Downloads/research/tse/GNNRecovery/example/dataset/Bash/include")
